"""
Interface definitions for continual learning models.

This module defines the core interfaces that all model implementations must follow,
establishing a consistent API for continual learning models.
"""

from abc import abstractmethod
from typing import List, Optional

import torch
import torch.nn as nn

from continuallearning.interfaces.core.component import (
    ComponentInterface,
)
from continuallearning.interfaces.types import ModelOutput
from continuallearning.interfaces.models.arch.backbone import BackboneInterface
from continuallearning.interfaces.core.task_related_interface import (
    TaskAwareInterface,
    TaskIrrelevantInterface,
)
from continuallearning.interfaces.models.pefts.adapter import (
    HookBasedAdapterInterface,
    AdapterInterface,
)
from continuallearning.interfaces.models.arch.head import (
    HeadInterface,
    MultiHeadInterface,
)


class ModelInterface(
    ComponentInterface, nn.Module, TaskIrrelevantInterface, TaskAwareInterface
):
    """
    Interface for continual learning models.

    This interface defines the API that all continual learning model
    implementations must follow, ensuring consistent behavior.
    """

    backbone: BackboneInterface
    heads: MultiHeadInterface
    adapters: HookBasedAdapterInterface

    @abstractmethod
    def forward(
        self, x: torch.Tensor, task_ids: Optional[List[int]] = None, **kwargs
    ) -> ModelOutput:
        """
        Forward pass through the model.

        Args:
            x: Input tensor
            task_ids: Optional list of task identifiers
            **kwargs: Additional arguments

        Returns:
            Union[torch.Tensor, ModelOutput]: Model outputs, either logits or a structured output
        """
        pass

    @abstractmethod
    def get_features(
        self, x: torch.Tensor, task_ids: Optional[List[int]] = None, **kwargs
    ) -> torch.Tensor:
        """
        Extract features without applying the classification head.

        Args:
            x: Input tensor
            task_ids: Optional list of task identifiers
            **kwargs: Additional arguments

        Returns:
            torch.Tensor: Feature representations
        """
        pass

    @abstractmethod
    def prepare_task(self, task_id: int, **kwargs) -> None:
        pass
